#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MongoDB PCAP分析工具
此脚本用于解析PCAP文件中的MongoDB协议数据包并显示交互信息

功能:
- 解析PCAP捕获文件中的MongoDB协议数据包
- 分析并显示MongoDB客户端和服务器之间的通信
- 支持多种MongoDB操作码和BSON数据类型的解析
- 按TCP流组织和显示数据包

使用方法:
python mongo_pcap_analyzer.py <pcap文件路径>

依赖:
- pyshark: 用于读取和解析PCAP文件
- 其他标准库: struct, binascii, datetime等
"""

import sys
import struct
import binascii
from collections import defaultdict
import datetime
import logging
import pyshark
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,  # 恢复为INFO级别
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('mongo_pcap_analyzer')

# MongoDB操作码定义
OP_CODES = {
    1: "OP_REPLY",         # 服务器对查询的响应
    1000: "OP_MESSAGE",    # 已废弃
    2001: "OP_UPDATE",     # 更新文档
    2002: "OP_INSERT",     # 插入文档
    2004: "OP_QUERY",      # 查询文档
    2005: "OP_GET_MORE",   # 获取更多查询结果
    2006: "OP_DELETE",     # 删除文档
    2007: "OP_KILL_CURSORS", # 终止游标
    2010: "OP_COMMAND",    # 数据库命令 (MongoDB 3.6+)
    2011: "OP_COMMANDREPLY", # 命令响应
    2012: "OP_COMPRESSED", # 压缩消息
    2013: "OP_MSG"         # 消息 (MongoDB 3.6+)
}

# BSON类型代码
BSON_TYPES = {
    0x01: "Double",       # 8字节双精度浮点数
    0x02: "String",       # UTF-8字符串
    0x03: "Document",     # 嵌入式文档
    0x04: "Array",        # 数组
    0x05: "Binary",       # 二进制数据
    0x07: "ObjectId",     # MongoDB对象ID (12字节)
    0x08: "Boolean",      # 布尔值
    0x09: "UTC Datetime", # UTC日期时间 (8字节时间戳)
    0x0A: "Null",         # 空值
    0x10: "Int32",        # 32位整数
    0x12: "Int64"         # 64位整数
}

# 压缩算法
COMPRESSORS = {
    0: "NOOP",            # 无压缩
    1: "Snappy",          # Snappy压缩
    2: "Zlib",            # Zlib压缩
    3: "Zstd"             # Zstandard压缩
}

class MongoPacketAnalyzer:
    """
    MongoDB数据包分析器
    
    此类用于解析PCAP文件中的MongoDB协议数据包，并按TCP流组织数据
    """
    def __init__(self, debug=False):  # 关闭默认调试模式
        """
        初始化MongoDB数据包分析器
        
        参数:
            debug (bool): 是否启用调试模式
        """
        self.streams = defaultdict(list)  # 按TCP流索引存储数据包
        self.debug = debug
        self.packet_count = 0
        self.mongo_packet_count = 0
        
    def parse_pcap(self, pcap_file):
        """
        解析PCAP文件并提取MongoDB数据包
        
        参数:
            pcap_file (str): PCAP文件路径
            
        返回:
            dict: 按TCP流索引组织的MongoDB数据包字典
        """
        logger.info(f"正在解析PCAP文件: {pcap_file}")
        
        # 使用pyshark读取pcap文件
        try:
            capture = pyshark.FileCapture(pcap_file, display_filter="tcp")
            logger.info("成功打开PCAP文件")
        except Exception as e:
            logger.error(f"打开PCAP文件失败: {e}")
            raise
        
        for packet in capture:
            self.packet_count += 1
            
            if self.packet_count % 1000 == 0:
                logger.info(f"已处理 {self.packet_count} 个数据包...")
                
            try:
                if 'TCP' in packet:
                    logger.debug(f"发现TCP数据包 #{self.packet_count}")
                    
                    # 获取TCP流索引和负载数据
                    stream_index = int(packet.tcp.stream)
                    src_port = packet.tcp.srcport
                    dst_port = packet.tcp.dstport
                    direction = f"{packet.ip.src}:{src_port} -> {packet.ip.dst}:{dst_port}"
                    
                    # 尝试获取原始数据
                    raw_data = None
                    if hasattr(packet.tcp, 'payload'):
                        logger.debug(f"数据包 #{self.packet_count} 含有TCP负载")
                        hex_data = packet.tcp.payload.replace(':', '')
                        
                        try:
                            raw_data = binascii.unhexlify(hex_data)
                            logger.debug(f"负载长度: {len(raw_data)} 字节")
                        except binascii.Error as e:
                            logger.warning(f"解析TCP负载失败: {e}, 数据: {hex_data[:20]}...")
                            continue
                            
                        if raw_data and len(raw_data) > 16:  # 至少要有MongoDB消息头(16字节)
                            logger.debug(f"TCP负载长度足够构成MongoDB消息头: {len(raw_data)} 字节")
                            # 尝试解析MongoDB数据包头部以快速检查合法性
                            try:
                                message_length, request_id, response_to, op_code = struct.unpack("<IIII", raw_data[:16])
                                logger.debug(f"尝试解析MongoDB头: 长度={message_length}, 请求ID={request_id}, 响应ID={response_to}, 操作码={op_code}")
                                
                                # 检查操作码是否在已知MongoDB操作码列表中
                                if op_code not in OP_CODES:
                                    logger.debug(f"未知MongoDB操作码: {op_code}，跳过此数据包")
                                    continue
                                    
                                # 检查消息长度是否合理
                                if message_length < 16 or message_length > 1000000:  # 设置一个合理的上限
                                    logger.debug(f"MongoDB消息长度不合理: {message_length}，跳过此数据包")
                                    continue
                            except struct.error:
                                logger.debug("无法解析MongoDB消息头，跳过此数据包")
                                continue
                                
                            timestamp = float(packet.sniff_timestamp)
                            parsed_data = self.parse_mongo_packet(raw_data, direction, timestamp)
                            if parsed_data:
                                logger.debug(f"成功解析MongoDB数据包: {parsed_data['op_name']}")
                                self.mongo_packet_count += 1
                                self.streams[stream_index].append(parsed_data)
                            else:
                                logger.debug("parse_mongo_packet返回None，无法识别为MongoDB数据包")
                    else:
                        logger.debug(f"数据包 #{self.packet_count} 没有TCP负载")
            except Exception as e:
                logger.error(f"解析数据包时出错(包 #{self.packet_count}): {e}")
                if self.debug:
                    import traceback
                    logger.debug(traceback.format_exc())
                continue
                
        capture.close()
        logger.info(f"PCAP解析完成。总计处理 {self.packet_count} 个数据包，识别出 {self.mongo_packet_count} 个MongoDB数据包。")
        return self.streams
    
    def parse_mongo_packet(self, data, direction, timestamp):
        """
        解析MongoDB数据包
        
        参数:
            data (bytes): 原始数据包负载
            direction (str): 数据包方向 (源IP:端口 -> 目标IP:端口)
            timestamp (float): 数据包时间戳
            
        返回:
            dict: 解析后的MongoDB数据包信息
        """
        if len(data) < 16:
            if self.debug:
                logger.debug(f"数据包太小，无法构成MongoDB消息: {len(data)} 字节")
            return None
            
        # 解析消息头 (16字节)
        try:
            message_length, request_id, response_to, op_code = struct.unpack("<IIII", data[:16])
        except struct.error as e:
            logger.warning(f"解析消息头失败: {e}")
            return None
        
        if message_length > len(data):
            logger.warning(f"警告: 消息长度({message_length})大于数据包长度({len(data)})，可能是分片数据包")
            return None
            
        op_name = OP_CODES.get(op_code, f"未知操作码({op_code})")
        
        if self.debug:
            logger.debug(f"发现MongoDB数据包: {op_name}, 请求ID: {request_id}, 长度: {message_length}")
        
        packet_info = {
            "timestamp": datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S.%f'),
            "direction": direction,
            "message_length": message_length,
            "request_id": request_id,
            "response_to": response_to,
            "op_code": op_code,
            "op_name": op_name,
            "body": {}
        }
        
        # 解析消息体 (根据操作码)
        if len(data) > 16:
            body_data = data[16:message_length]
            offset = 0
            
            try:
                if op_code == 2013:  # OP_MSG
                    packet_info["body"] = self.parse_op_msg(body_data)
                elif op_code == 2012:  # OP_COMPRESSED
                    packet_info["body"] = self.parse_op_compressed(body_data)
                elif op_code == 2004:  # OP_QUERY
                    packet_info["body"] = self.parse_op_query(body_data)
                elif op_code == 1:  # OP_REPLY
                    packet_info["body"] = self.parse_op_reply(body_data)
                elif op_code == 2001:  # OP_UPDATE
                    packet_info["body"] = self.parse_op_update(body_data)
                elif op_code == 2002:  # OP_INSERT
                    packet_info["body"] = self.parse_op_insert(body_data)
                elif op_code == 2006:  # OP_DELETE
                    packet_info["body"] = self.parse_op_delete(body_data)
                elif op_code == 2005:  # OP_GET_MORE
                    packet_info["body"] = {"raw": "OP_GET_MORE解析未实现"}
                elif op_code == 2007:  # OP_KILL_CURSORS
                    packet_info["body"] = {"raw": "OP_KILL_CURSORS解析未实现"}
                else:
                    packet_info["body"] = {"raw": f"未实现的操作码: {op_code}", "hex_data": binascii.hexlify(body_data[:50]).decode('utf-8') + ("..." if len(body_data) > 50 else "")}
            except Exception as e:
                logger.warning(f"解析消息体时出错(操作码 {op_name}): {e}")
                if self.debug:
                    import traceback
                    logger.debug(traceback.format_exc())
                packet_info["body"] = {"error": f"解析失败: {str(e)}", "raw_hex": binascii.hexlify(body_data[:50]).decode('utf-8') + ("..." if len(body_data) > 50 else "")}
        
        return packet_info
        
    def parse_op_msg(self, data):
        """
        解析OP_MSG消息体 (MongoDB 3.6+)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_MSG数据
        """
        if len(data) < 4:
            return {"error": "OP_MSG数据不完整"}
            
        flag_bits = struct.unpack("<I", data[:4])[0]
        sections = []
        offset = 4
        
        # 检查flagBits
        checksum_present = bool(flag_bits & 0x00000001)  # 是否存在校验和
        more_to_come = bool(flag_bits & 0x00000002)      # 是否有更多消息
        exhaust_allowed = bool(flag_bits & 0x00010000)   # 是否允许耗尽游标
        
        # 解析sections
        while offset < len(data) - (4 if checksum_present else 0):
            if offset >= len(data):
                break
                
            kind = data[offset]
            offset += 1
            
            if kind == 0:  # Body section (单个BSON文档)
                doc, doc_size = self.parse_bson(data[offset:])
                if doc:
                    sections.append({"kind": 0, "document": doc})
                    offset += doc_size
                else:
                    if self.debug:
                        logger.debug(f"无法解析OP_MSG Body section中的BSON文档，位置: {offset}")
                    break
            elif kind == 1:  # Document sequence (多个BSON文档序列)
                if offset + 4 >= len(data):
                    break
                    
                size = struct.unpack("<I", data[offset:offset+4])[0]
                offset += 4
                
                # 解析identifier (C风格字符串)
                identifier_end = data.find(b'\x00', offset)
                if identifier_end == -1:
                    if self.debug:
                        logger.debug(f"无法在OP_MSG Document sequence中找到identifier结束符，位置: {offset}")
                    break
                    
                identifier = data[offset:identifier_end].decode('utf-8', errors='replace')
                offset = identifier_end + 1
                
                # 解析documents
                documents = []
                sequence_end = offset + size - (identifier_end - offset + 1) - 4
                
                while offset < sequence_end:
                    doc, doc_size = self.parse_bson(data[offset:])
                    if doc:
                        documents.append(doc)
                        offset += doc_size
                    else:
                        if self.debug:
                            logger.debug(f"无法解析OP_MSG Document sequence中的BSON文档，位置: {offset}")
                        break
                        
                sections.append({
                    "kind": 1,
                    "identifier": identifier,
                    "documents": documents
                })
            else:
                if self.debug:
                    logger.debug(f"未知的OP_MSG section类型: {kind}，位置: {offset}")
                break
        
        # 解析校验和
        checksum = None
        if checksum_present and offset + 4 <= len(data):
            checksum = struct.unpack("<I", data[offset:offset+4])[0]
        
        return {
            "flag_bits": flag_bits,
            "checksum_present": checksum_present,
            "more_to_come": more_to_come,
            "exhaust_allowed": exhaust_allowed,
            "sections": sections,
            "checksum": checksum
        }
    
    def parse_op_compressed(self, data):
        """
        解析OP_COMPRESSED消息体
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_COMPRESSED数据
        """
        if len(data) < 9:
            return {"error": "OP_COMPRESSED数据不完整"}
            
        original_opcode, uncompressed_size, compressor_id = struct.unpack("<IIB", data[:9])
        compressed_data = data[9:]
        
        original_op_name = OP_CODES.get(original_opcode, f"未知操作码({original_opcode})")
        compressor_name = COMPRESSORS.get(compressor_id, f"未知压缩算法({compressor_id})")
        
        if self.debug:
            logger.debug(f"OP_COMPRESSED: 原始操作码={original_op_name}, 压缩算法={compressor_name}, 解压后大小={uncompressed_size}")
        
        return {
            "original_opcode": original_opcode,
            "original_opcode_name": original_op_name,
            "uncompressed_size": uncompressed_size,
            "compressor_id": compressor_id,
            "compressor_name": compressor_name,
            "compressed_data_size": len(compressed_data),
            # 注意: 实际解压代码需要相应的压缩库支持
            "compressed_data_hex": binascii.hexlify(compressed_data[:20]).decode('utf-8') + "..." if len(compressed_data) > 20 else binascii.hexlify(compressed_data).decode('utf-8')
        }
    
    def parse_op_query(self, data):
        """
        解析OP_QUERY消息体 (数据库查询)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_QUERY数据
        """
        if len(data) < 12:
            return {"error": "OP_QUERY数据不完整"}
            
        flags = struct.unpack("<I", data[:4])[0]
        offset = 4
        
        # 解析查询标志
        query_flags = {
            "tailable_cursor": bool(flags & 0x02),       # 可追踪游标
            "slave_ok": bool(flags & 0x04),             # 允许从从节点读取
            "oplog_replay": bool(flags & 0x08),         # 用于复制操作日志
            "no_cursor_timeout": bool(flags & 0x10),    # 游标不超时
            "await_data": bool(flags & 0x20),           # 等待数据
            "exhaust": bool(flags & 0x40),              # 耗尽游标
            "partial": bool(flags & 0x80)               # 允许部分结果
        }
        
        # 解析集合名称 (C风格字符串)
        coll_name_end = data.find(b'\x00', offset)
        if coll_name_end == -1:
            return {"error": "无法解析集合名称"}
            
        collection_name = data[offset:coll_name_end].decode('utf-8', errors='replace')
        offset = coll_name_end + 1
        
        number_to_skip, number_to_return = struct.unpack("<II", data[offset:offset+8])
        offset += 8
        
        # 解析查询文档
        query_doc, query_size = self.parse_bson(data[offset:])
        if not query_doc:
            if self.debug:
                logger.debug(f"无法解析OP_QUERY中的查询文档，位置: {offset}")
            query_doc = {"error": "无法解析查询文档"}
            query_size = 0
            
        offset += query_size
        
        # 可选的投影文档
        projection_doc = None
        if offset < len(data):
            projection_doc, _ = self.parse_bson(data[offset:])
            if not projection_doc and self.debug:
                logger.debug(f"无法解析OP_QUERY中的投影文档，位置: {offset}")
        
        if self.debug:
            logger.debug(f"OP_QUERY: 集合={collection_name}, 跳过={number_to_skip}, 返回={number_to_return}")
        
        return {
            "flags": flags,
            "flag_details": query_flags,
            "collection_name": collection_name,
            "number_to_skip": number_to_skip,
            "number_to_return": number_to_return,
            "query": query_doc,
            "projection": projection_doc
        }
    
    def parse_op_reply(self, data):
        """
        解析OP_REPLY消息体 (服务器响应)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_REPLY数据
        """
        if len(data) < 20:
            return {"error": "OP_REPLY数据不完整"}
            
        response_flags, cursor_id, starting_from, number_returned = struct.unpack("<IQii", data[:20])
        offset = 20
        
        # 解析响应标志
        reply_flags = {
            "cursor_not_found": bool(response_flags & 0x01),   # 游标未找到
            "query_failure": bool(response_flags & 0x02),      # 查询失败
            "shard_config_stale": bool(response_flags & 0x04), # 分片配置过期
            "await_capable": bool(response_flags & 0x08)       # 支持等待数据
        }
        
        # 解析返回的文档
        documents = []
        docs_parsed = 0
        
        for i in range(number_returned):
            if offset >= len(data):
                if self.debug:
                    logger.warning(f"OP_REPLY声明返回{number_returned}个文档，但只能解析{docs_parsed}个")
                break
                
            doc, doc_size = self.parse_bson(data[offset:])
            if doc:
                documents.append(doc)
                offset += doc_size
                docs_parsed += 1
            else:
                if self.debug:
                    logger.debug(f"无法解析OP_REPLY中的文档 #{i+1}，位置: {offset}")
                break
        
        if self.debug:
            logger.debug(f"OP_REPLY: 游标ID={cursor_id}, 返回文档数={number_returned}, 实际解析={docs_parsed}")
        
        return {
            "response_flags": response_flags,
            "flag_details": reply_flags,
            "cursor_id": cursor_id,
            "starting_from": starting_from,
            "number_returned": number_returned,
            "documents": documents
        }
    
    def parse_op_update(self, data):
        """
        解析OP_UPDATE消息体 (更新操作)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_UPDATE数据
        """
        if len(data) < 8:
            return {"error": "OP_UPDATE数据不完整"}
            
        zero = struct.unpack("<I", data[:4])[0]  # 保留字段，总是0
        if zero != 0 and self.debug:
            logger.warning(f"OP_UPDATE保留字段不为0: {zero}")
            
        offset = 4
        
        # 解析集合名称 (C风格字符串)
        coll_name_end = data.find(b'\x00', offset)
        if coll_name_end == -1:
            return {"error": "无法解析集合名称"}
            
        collection_name = data[offset:coll_name_end].decode('utf-8', errors='replace')
        offset = coll_name_end + 1
        
        flags = struct.unpack("<I", data[offset:offset+4])[0]
        offset += 4
        
        # 解析标志
        upsert = bool(flags & 1)        # 如果不存在则插入
        multi_update = bool(flags & 2)  # 更新多个文档
        
        # 解析选择器和更新文档
        selector_doc, selector_size = self.parse_bson(data[offset:])
        if not selector_doc:
            if self.debug:
                logger.debug(f"无法解析OP_UPDATE中的选择器文档，位置: {offset}")
            selector_doc = {"error": "无法解析选择器文档"}
            selector_size = 0
            
        offset += selector_size
        
        update_doc = None
        if offset < len(data):
            update_doc, _ = self.parse_bson(data[offset:])
            if not update_doc and self.debug:
                logger.debug(f"无法解析OP_UPDATE中的更新文档，位置: {offset}")
        
        if self.debug:
            logger.debug(f"OP_UPDATE: 集合={collection_name}, upsert={upsert}, multi={multi_update}")
        
        return {
            "collection_name": collection_name,
            "flags": flags,
            "upsert": upsert,
            "multi_update": multi_update,
            "selector": selector_doc,
            "update": update_doc
        }
    
    def parse_op_insert(self, data):
        """
        解析OP_INSERT消息体 (插入操作)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_INSERT数据
        """
        if len(data) < 4:
            return {"error": "OP_INSERT数据不完整"}
            
        flags = struct.unpack("<I", data[:4])[0]
        offset = 4
        
        # 解析标志
        continue_on_error = bool(flags & 1)  # 出错时继续处理
        
        # 解析集合名称 (C风格字符串)
        coll_name_end = data.find(b'\x00', offset)
        if coll_name_end == -1:
            return {"error": "无法解析集合名称"}
            
        collection_name = data[offset:coll_name_end].decode('utf-8', errors='replace')
        offset = coll_name_end + 1
        
        # 解析要插入的文档
        documents = []
        doc_count = 0
        
        while offset < len(data):
            doc, doc_size = self.parse_bson(data[offset:])
            if doc:
                documents.append(doc)
                offset += doc_size
                doc_count += 1
            else:
                if self.debug:
                    logger.debug(f"无法解析OP_INSERT中的文档 #{doc_count+1}，位置: {offset}")
                break
        
        if self.debug:
            logger.debug(f"OP_INSERT: 集合={collection_name}, 文档数={doc_count}, continue_on_error={continue_on_error}")
        
        return {
            "flags": flags,
            "continue_on_error": continue_on_error,
            "collection_name": collection_name,
            "documents": documents
        }
    
    def parse_op_delete(self, data):
        """
        解析OP_DELETE消息体 (删除操作)
        
        参数:
            data (bytes): 消息体数据
            
        返回:
            dict: 解析后的OP_DELETE数据
        """
        if len(data) < 8:
            return {"error": "OP_DELETE数据不完整"}
            
        zero = struct.unpack("<I", data[:4])[0]  # 保留字段，总是0
        if zero != 0 and self.debug:
            logger.warning(f"OP_DELETE保留字段不为0: {zero}")
            
        offset = 4
        
        # 解析集合名称 (C风格字符串)
        coll_name_end = data.find(b'\x00', offset)
        if coll_name_end == -1:
            return {"error": "无法解析集合名称"}
            
        collection_name = data[offset:coll_name_end].decode('utf-8', errors='replace')
        offset = coll_name_end + 1
        
        flags = struct.unpack("<I", data[offset:offset+4])[0]
        offset += 4
        
        # 解析标志
        single_remove = bool(flags & 1)  # 只删除一个文档
        
        # 解析选择器文档
        selector_doc = None
        if offset < len(data):
            selector_doc, _ = self.parse_bson(data[offset:])
            if not selector_doc and self.debug:
                logger.debug(f"无法解析OP_DELETE中的选择器文档，位置: {offset}")
        
        if self.debug:
            logger.debug(f"OP_DELETE: 集合={collection_name}, single_remove={single_remove}")
        
        return {
            "collection_name": collection_name,
            "flags": flags,
            "single_remove": single_remove,
            "selector": selector_doc
        }
    
    def parse_bson(self, data):
        """
        解析BSON文档并返回(文档, 文档大小)元组
        
        参数:
            data (bytes): 包含BSON文档的二进制数据
            
        返回:
            tuple: (解析后的文档字典, 文档大小) 或 (None, 0)表示失败
        """
        if len(data) < 4:
            if self.debug:
                logger.debug("BSON数据太小，无法包含长度前缀")
            return None, 0
            
        try:
            # 获取BSON文档长度
            doc_size = struct.unpack("<i", data[:4])[0]
            
            if doc_size <= 0:
                if self.debug:
                    logger.warning(f"BSON文档长度无效: {doc_size}")
                return None, 0
                
            if doc_size > len(data):
                if self.debug:
                    logger.warning(f"BSON文档长度({doc_size})大于可用数据长度({len(data)})")
                return None, 0
                
            # 从二进制数据中提取BSON文档
            bson_data = data[:doc_size]
            
            # 解析BSON文档元素
            document = {}
            offset = 4
            
            while offset < doc_size - 1:  # -1排除文档末尾的\x00
                # 获取元素类型
                element_type = bson_data[offset]
                offset += 1
                
                # 如果是文档结束标记(0)，跳出循环
                if element_type == 0:
                    break
                    
                # 获取元素名称 (C风格字符串)
                name_end = bson_data.find(b'\x00', offset)
                if name_end == -1:
                    if self.debug:
                        logger.warning(f"无法在BSON中找到元素名称的结束符，位置: {offset}")
                    break
                
                name = bson_data[offset:name_end].decode('utf-8', errors='replace')
                offset = name_end + 1
                
                # 根据类型解析值
                value = self.parse_bson_element(element_type, bson_data, offset)
                if value is None:
                    if self.debug:
                        logger.warning(f"无法解析BSON元素 '{name}' (类型: 0x{element_type:02x})，位置: {offset}")
                    break
                    
                document[name] = value['value']
                offset += value['size']
            
            # 检查文档是否以\x00结束
            if offset >= doc_size or bson_data[doc_size-1] != 0:
                if self.debug and offset < doc_size:
                    logger.warning(f"BSON文档不是以\\x00结束的: 0x{bson_data[doc_size-1]:02x}")
            
            return document, doc_size
        except Exception as e:
            logger.warning(f"解析BSON文档时出错: {e}")
            if self.debug:
                import traceback
                logger.debug(traceback.format_exc())
            return None, 0
    
    def parse_bson_element(self, element_type, data, offset):
        """
        解析BSON元素并返回{value: 值, size: 占用字节数}
        
        参数:
            element_type (int): BSON元素类型代码
            data (bytes): 包含BSON文档的二进制数据
            offset (int): 元素值在数据中的起始位置
            
        返回:
            dict: 包含解析值和大小的字典，或None表示失败
        """
        type_name = BSON_TYPES.get(element_type, f"未知类型(0x{element_type:02x})")
        
        try:
            if element_type == 0x01:  # Double
                value = struct.unpack("<d", data[offset:offset+8])[0]
                return {"value": value, "size": 8}
                
            elif element_type == 0x02:  # String
                if offset + 4 > len(data):
                    if self.debug:
                        logger.debug(f"BSON String: 数据不足以包含字符串长度，位置: {offset}")
                    return None
                    
                str_size = struct.unpack("<i", data[offset:offset+4])[0]
                if str_size <= 0:
                    if self.debug:
                        logger.warning(f"BSON String: 无效字符串长度 {str_size}")
                    return None
                    
                offset += 4
                if offset + str_size > len(data):
                    if self.debug:
                        logger.debug(f"BSON String: 数据不足以包含完整字符串(长度 {str_size})，位置: {offset}")
                    return None
                    
                # 注意: str_size包括结尾的NULL字节
                value = data[offset:offset+str_size-1].decode('utf-8', errors='replace')
                return {"value": value, "size": 4 + str_size}
                
            elif element_type == 0x03 or element_type == 0x04:  # Document or Array
                embedded_doc, embedded_size = self.parse_bson(data[offset:])
                if embedded_doc is None:
                    return None
                    
                # 如果是数组，将文档转换为Python列表
                if element_type == 0x04:
                    # 数组在BSON中存储为特殊格式的文档，键名为"0", "1", "2"等
                    try:
                        max_index = max(int(k) for k in embedded_doc.keys() if k.isdigit())
                        array = [None] * (max_index + 1)
                        for k, v in embedded_doc.items():
                            if k.isdigit():
                                array[int(k)] = v
                        embedded_doc = array
                    except (ValueError, IndexError) as e:
                        if self.debug:
                            logger.warning(f"无法将BSON文档转换为数组: {e}")
                
                return {"value": embedded_doc, "size": embedded_size}
                
            elif element_type == 0x05:  # Binary
                if offset + 5 > len(data):
                    if self.debug:
                        logger.debug(f"BSON Binary: 数据不足以包含二进制数据长度和子类型，位置: {offset}")
                    return None
                    
                bin_size = struct.unpack("<i", data[offset:offset+4])[0]
                if bin_size < 0:
                    if self.debug:
                        logger.warning(f"BSON Binary: 无效二进制长度 {bin_size}")
                    return None
                    
                offset += 4
                subtype = data[offset]
                offset += 1
                
                if offset + bin_size > len(data):
                    if self.debug:
                        logger.debug(f"BSON Binary: 数据不足以包含完整二进制数据(长度 {bin_size})，位置: {offset}")
                    return None
                    
                binary_data = data[offset:offset+bin_size]
                
                # 定义二进制子类型
                subtype_names = {
                    0x00: "通用二进制",
                    0x01: "函数",
                    0x02: "旧式二进制(弃用)",
                    0x03: "UUID(旧式)",
                    0x04: "UUID",
                    0x05: "MD5",
                    0x80: "用户自定义"
                }
                subtype_name = subtype_names.get(subtype, f"未知子类型(0x{subtype:02x})")
                
                # 如果是UUID，尝试格式化显示
                uuid_str = None
                if subtype in (0x03, 0x04) and len(binary_data) == 16:
                    try:
                        import uuid
                        uuid_obj = uuid.UUID(bytes=binary_data)
                        uuid_str = str(uuid_obj)
                    except:
                        pass
                
                value = {
                    "subtype": subtype,
                    "subtype_name": subtype_name,
                    "size": bin_size,
                    "data": binascii.hexlify(binary_data).decode('utf-8')
                }
                
                if uuid_str:
                    value["uuid"] = uuid_str
                    
                return {"value": value, "size": 5 + bin_size}
                
            elif element_type == 0x07:  # ObjectId
                if offset + 12 > len(data):
                    if self.debug:
                        logger.debug(f"BSON ObjectId: 数据不足以包含完整ObjectId(需要12字节)，位置: {offset}")
                    return None
                    
                object_id = binascii.hexlify(data[offset:offset+12]).decode('utf-8')
                
                # 从ObjectId中提取时间戳（前4字节）
                timestamp = struct.unpack(">I", data[offset:offset+4])[0]
                creation_time = datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
                
                return {
                    "value": {
                        "id": object_id,
                        "timestamp": timestamp,
                        "creation_time": creation_time
                    }, 
                    "size": 12
                }
                
            elif element_type == 0x08:  # Boolean
                if offset >= len(data):
                    if self.debug:
                        logger.debug(f"BSON Boolean: 数据不足，位置: {offset}")
                    return None
                    
                value = bool(data[offset])
                return {"value": value, "size": 1}
                
            elif element_type == 0x09:  # UTC Datetime
                if offset + 8 > len(data):
                    if self.debug:
                        logger.debug(f"BSON Datetime: 数据不足以包含完整时间戳(需要8字节)，位置: {offset}")
                    return None
                    
                value = struct.unpack("<q", data[offset:offset+8])[0]
                dt = datetime.datetime.fromtimestamp(value/1000.0).strftime('%Y-%m-%d %H:%M:%S.%f')
                return {"value": {"timestamp": value, "datetime": dt}, "size": 8}
                
            elif element_type == 0x0A:  # Null
                return {"value": None, "size": 0}
                
            elif element_type == 0x10:  # Int32
                if offset + 4 > len(data):
                    if self.debug:
                        logger.debug(f"BSON Int32: 数据不足以包含完整整数(需要4字节)，位置: {offset}")
                    return None
                    
                value = struct.unpack("<i", data[offset:offset+4])[0]
                return {"value": value, "size": 4}
                
            elif element_type == 0x12:  # Int64
                if offset + 8 > len(data):
                    if self.debug:
                        logger.debug(f"BSON Int64: 数据不足以包含完整整数(需要8字节)，位置: {offset}")
                    return None
                    
                value = struct.unpack("<q", data[offset:offset+8])[0]
                return {"value": value, "size": 8}
                
            else:
                if self.debug:
                    logger.warning(f"未支持的BSON类型: 0x{element_type:02x}")
                return {"value": f"未支持的BSON类型: 0x{element_type:02x}", "size": 0}
        except Exception as e:
            logger.warning(f"解析BSON元素(类型: {type_name})时出错: {e}")
            if self.debug:
                import traceback
                logger.debug(traceback.format_exc())
            return None

def print_mongo_packet(packet):
    """
    格式化输出MongoDB数据包信息
    
    参数:
        packet (dict): 解析后的MongoDB数据包信息
    """
    print(f"时间戳: {packet['timestamp']}")
    print(f"方向: {packet['direction']}")
    print(f"消息长度: {packet['message_length']} 字节")
    print(f"请求ID: {packet['request_id']}")
    print(f"响应ID: {packet['response_to']}")
    print(f"操作码: {packet['op_code']} ({packet['op_name']})")
    
    # 根据操作类型格式化输出消息体
    print("消息体:")
    pretty_print_dict(packet['body'], 1)
    print("-" * 80)

def pretty_print_dict(d, indent=0):
    """
    美化输出字典
    
    参数:
        d (dict): 要输出的字典
        indent (int): 缩进级别
    """
    for key, value in d.items():
        indent_str = "  " * indent
        if isinstance(value, dict):
            print(f"{indent_str}{key}:")
            pretty_print_dict(value, indent + 1)
        elif isinstance(value, list):
            if len(value) == 0:
                print(f"{indent_str}{key}: []")
            else:
                print(f"{indent_str}{key}: [")
                for i, item in enumerate(value):
                    if isinstance(item, dict):
                        print(f"{indent_str}  项目 {i}:")
                        pretty_print_dict(item, indent + 2)
                    else:
                        print(f"{indent_str}  项目 {i}: {item}")
                print(f"{indent_str}]")
        else:
            print(f"{indent_str}{key}: {value}")

def main():
    """主函数 - 程序入口点"""
    
    try:
        # 检查参数
        if len(sys.argv) < 2:
            print("用法: python mongo_pcap_analyzer.py <pcap文件路径>")
            sys.exit(1)
        
        pcap_file = sys.argv[1]
        
        if not os.path.exists(pcap_file):
            print(f"错误: PCAP文件不存在: {pcap_file}")
            sys.exit(1)
            
        # 打印欢迎信息
        print(f"\n=== MongoDB PCAP 分析工具 ===")
        print(f"分析文件: {pcap_file}")
        print(f"开始时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("-" * 80)
        
        # 实例化分析器并运行
        start_time = datetime.datetime.now()
        analyzer = MongoPacketAnalyzer(debug=False)  # 关闭调试模式
        streams = analyzer.parse_pcap(pcap_file)
        end_time = datetime.datetime.now()
        duration = (end_time - start_time).total_seconds()
        
        # 打印分析摘要
        print("\n=== 分析摘要 ===")
        print(f"分析耗时: {duration:.2f} 秒")
        print(f"发现 {len(streams)} 个TCP流")
        print(f"共解析 {analyzer.mongo_packet_count} 个MongoDB数据包")
        print("-" * 80)
        
        # 输出每个流中的MongoDB数据包
        if len(streams) > 0:
            tcp_stream_count = 0
            for stream_id, packets in streams.items():
                if len(packets) > 0:
                    tcp_stream_count += 1
                    print(f"\n=== TCP 流 #{stream_id} ===")
                    print(f"数据包数量: {len(packets)}")
                    
                    # 提取流的基本信息
                    first_packet = packets[0]
                    clients = set()
                    servers = set()
                    
                    for p in packets:
                        src, dst = p['direction'].split(' -> ')
                        if 'query' in p.get('body', {}) or 'selector' in p.get('body', {}):
                            clients.add(src)
                            servers.add(dst)
                        elif 'documents' in p.get('body', {}):
                            servers.add(src)
                            clients.add(dst)
                    
                    if clients:
                        print(f"可能的客户端: {', '.join(clients)}")
                    if servers:
                        print(f"可能的服务器: {', '.join(servers)}")
                    
                    # 计算操作类型统计
                    op_stats = {}
                    for p in packets:
                        op_name = p['op_name']
                        op_stats[op_name] = op_stats.get(op_name, 0) + 1
                    
                    print("操作统计:")
                    for op, count in op_stats.items():
                        print(f"  {op}: {count} 次")
                
                    # 详细输出数据包
                    print("\n数据包详情:")
                    for i, packet in enumerate(packets):
                        print(f"\n包 #{i+1}:")
                        print_mongo_packet(packet)
                        
                        # 对于大量数据包，可以限制输出数量
                        if i >= 49 and len(packets) > 50:
                            print(f"\n... 还有 {len(packets) - 50} 个数据包未显示 ...")
                            break
            
            if tcp_stream_count == 0:
                print("\n未找到包含MongoDB数据包的TCP流。")
        else:
            print("\n未找到任何MongoDB数据包。")
            
        print("\n分析完成!")
    except KeyboardInterrupt:
        print("\n用户中断了分析过程")
        sys.exit(2)
    except Exception as e:
        print(f"\n处理PCAP文件时出错: {e}")
        import traceback
        print(traceback.format_exc())
        sys.exit(1)

if __name__ == "__main__":
    main() 
